package com.onewo.arthasdemo.handler;

import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.lang.instrument.Instrumentation;
import java.lang.instrument.UnmodifiableClassException;
import java.security.ProtectionDomain;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.function.Function;
import java.util.logging.Logger;
import java.util.stream.Collectors;

public class Handler implements ClassFileTransformer {

    public static Instrumentation inst;

    private static final Logger logger = Logger.getLogger(Handler.class.getName());

    private static ConcurrentMap<String,List<Function<byte[],byte[]>>> clzOpsMap = new ConcurrentHashMap<>();


    public void enhance(String clzNameAndOperation)  throws UnmodifiableClassException{
        try {
            final String[] split = clzNameAndOperation.split("#");
            String clzName = split[0];
            String opName = split[1];
            final List<Function<byte[], byte[]>> opsList = clzOpsMap.getOrDefault(clzName, new ArrayList<>());
            if (ClazzOps.clazzOpsMap.containsKey(opName)){
                opsList.add(ClazzOps.clazzOpsMap.get(opName));
            }
            clzOpsMap.put(clzName,opsList);

            final Class[] allLoadedClasses = inst.getAllLoadedClasses();
            final List<Class> collect = Arrays.stream(allLoadedClasses).filter(f -> f.getName().contains(clzName)).collect(Collectors.toList());
            for (Class c : collect) {
                logger.info("Start transform class: " + c.getSimpleName());
                inst.retransformClasses(c);//实际调用inst增强代码(JVMTI修改class)
                logger.info("Success to transform class: " + c);
            }
        } catch (Throwable t) {
            logger.warning("retransform failed.");
            if (t instanceof UnmodifiableClassException) {
                throw (UnmodifiableClassException) t;
            } else if (t instanceof RuntimeException) {
                throw (RuntimeException) t;
            } else {
                throw new RuntimeException(t);
            }
        }
    }

    @Override
    public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer){
        try {
            final String simpleName = classBeingRedefined.getSimpleName();
            final String[] actualName = simpleName.split("\\$\\$");//spring 代理类
            logger.info("start transform simpleClassName: "+ actualName[0]);
            if (clzOpsMap.containsKey(actualName[0])){
                logger.info("[transformd] transformd className: "+ className);
                final List<Function<byte[],byte[]>> ops = clzOpsMap.get(actualName[0]);
                byte[] tmp = classfileBuffer;
                for (Function<byte[], byte[]> op : ops) {
                    tmp = op.apply(tmp);
                }
                return tmp;
            }
            return classfileBuffer;
        } catch (Exception e) {
            return classfileBuffer;
        }
    }


    public static void setInst(Instrumentation inst) {
        Handler.inst = inst;
    }
}
